package com.zst.sw.server.pipeline;

import com.alibaba.fastjson2.JSON;
import com.zst.sw.server.module.TraceDataHandler;
import com.zst.sw.server.module.segment.Segment;
import com.zst.sw.server.module.segment.SegmentMapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.skywalking.apm.network.language.agent.v3.SegmentObject;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;
import reactor.core.scheduler.Schedulers;

import java.time.Duration;
import java.util.List;
import java.util.function.Consumer;

@Slf4j
public class TraceSegmentDataPipeline {
    private DataPipelineProperties prop;
    private TraceDataHandler dataHandler;
    private Consumer<SegmentObject> consumer;

    public TraceSegmentDataPipeline(DataPipelineProperties prop, TraceDataHandler dataHandler) {
        if (dataHandler == null) {
            throw new IllegalArgumentException();
        }
        if (prop == null) {
            prop = new DataPipelineProperties();
        }

        this.prop = prop;
        this.dataHandler = dataHandler;

        init();
    }

    public void init() {
        Flux.<SegmentObject>create(fluxSink -> consumer = fluxSink::next, FluxSink.OverflowStrategy.DROP)
                .onBackpressureDrop(segmentObject -> {
                    log.warn("trace segment data pipeline overflow, segmentObject={}", JSON.toJSONString(segmentObject));
                })
                .bufferTimeout(prop.getMaxBatchSize(), Duration.ofMillis(prop.getMaxBatchIntervalMs()))
                .publishOn(Schedulers.newSingle("trace-segment-data-pipeline"), prop.getMaxCacheBatchNum())
                .retry()
                .subscribe(this::handleSegmentObjects, throwable -> {
                    log.error("trace segment data pipeline error", throwable);
                });

    }

    public void publish(SegmentObject segmentObject) {
        if (segmentObject == null) {
            throw new IllegalArgumentException();
        }
        if (consumer == null) {
            throw new IllegalStateException("trace segment data pipeline not init");
        }

        consumer.accept(segmentObject);
    }

    private void handleSegmentObjects(List<SegmentObject> segmentObjects) {
        // TODO 添加Filter机制
        try {
            List<Segment> segments = SegmentMapper.map(segmentObjects);
            dataHandler.handle(segments);
        } catch (Exception e) {
            log.error("trace segment data pipeline handle Segment Object error", e);
        }
    }
}
